import os
import json
from PIL import Image # 用于获取图片尺寸

def yolo_to_coco(yolo_dataset_root, coco_output_root, class_names):
    """
    将YOLO格式数据集转换为COCO格式。

    Args:
        yolo_dataset_root (str): YOLO格式数据集的根目录，
                                 例如：包含 'train', 'val', 'test' 子目录。
        coco_output_root (str): 转换后的COCO格式数据集的输出根目录。
                                 例如：会在该目录下创建 'images' 和 'annotations' 目录。
        class_names (list): 一个字符串列表，按索引对应YOLO标注中的class_id。
                            例如：['垃圾袋', '塑料瓶', '易拉罐', ...]
    """

    os.makedirs(coco_output_root, exist_ok=True)

    # COCO JSON 结构模板
    coco_json_template = {
        "info": {
            "description": "Converted from YOLO format to COCO format",
            "version": "1.0",
            "year": 2025,
            "date_created": "2025-06-07"
        },
        "licenses": [], # 可选：根据您的许可信息添加
        "categories": [],
        "images": [],
        "annotations": []
    }

    # 填充 categories
    for i, name in enumerate(class_names):
        coco_json_template["categories"].append({
            "id": i,
            "name": name,
            "supercategory": "garbage" # 可选，根据您的需要更改
        })

    image_id_counter = 0
    annotation_id_counter = 0

    # 遍历 train, val, test 子目录
    for split in ['train', 'val', 'test']:
        yolo_images_dir = os.path.join(yolo_dataset_root, split, 'images')
        yolo_labels_dir = os.path.join(yolo_dataset_root, split, 'labels')
        
        # 为COCO格式创建对应的图片和标注目录
        coco_images_dir = os.path.join(coco_output_root, 'images', split)
        coco_annotations_dir = os.path.join(coco_output_root, 'annotations') # 标注文件通常集中存储

        os.makedirs(coco_images_dir, exist_ok=True)
        os.makedirs(coco_annotations_dir, exist_ok=True)
        
        # 每个split生成一个独立的COCO JSON文件
        split_coco_data = coco_json_template.copy()
        split_coco_data["images"] = []
        split_coco_data["annotations"] = []

        if not os.path.exists(yolo_images_dir):
            print(f"警告: 未找到 {split} 阶段的图片目录: {yolo_images_dir}")
            continue
        if not os.path.exists(yolo_labels_dir):
            print(f"警告: 未找到 {split} 阶段的标注目录: {yolo_labels_dir}")
            continue

        print(f"开始处理 {split} 数据...")
        for img_name in os.listdir(yolo_images_dir):
            if not (img_name.lower().endswith('.jpg') or img_name.lower().endswith('.png')):
                continue

            img_path = os.path.join(yolo_images_dir, img_name)
            label_name = os.path.splitext(img_name)[0] + '.txt'
            label_path = os.path.join(yolo_labels_dir, label_name)

            # 复制图片到新的COCO图片目录
            # 可以直接使用 os.link 或 shutil.copy2
            # 为了简单和兼容性，这里我们不复制图片，而是记录图片路径和ID
            # 如果需要物理复制图片，请取消注释下一行，并导入 shutil
            # import shutil
            # shutil.copy2(img_path, os.path.join(coco_images_dir, img_name))

            # 获取图片尺寸
            try:
                with Image.open(img_path) as img:
                    width, height = img.size
            except Exception as e:
                print(f"无法打开图片 {img_path}: {e}")
                continue

            current_image_id = image_id_counter
            image_id_counter += 1

            # 添加图片信息到COCO JSON
            split_coco_data["images"].append({
                "id": current_image_id,
                "width": width,
                "height": height,
                "file_name": os.path.join(split, img_name) # 相对路径，以便加载器可以找到图片
                # "file_name": img_name # 如果所有图片都在一个平面目录
            })

            # 解析YOLO标注文件
            if os.path.exists(label_path):
                with open(label_path, 'r') as f:
                    for line in f.readlines():
                        parts = line.strip().split(' ')
                        if len(parts) != 5:
                            print(f"警告: 标注文件 {label_path} 格式错误，跳过行: {line.strip()}")
                            continue

                        class_id = int(parts[0])
                        center_x_norm, center_y_norm, width_norm, height_norm = map(float, parts[1:])

                        # 将归一化YOLO坐标转换为COCO (x_min, y_min, width, height)
                        # 注意：YOLO的center_x, center_y, width, height都是归一化的
                        # COCO的bbox是(x_min, y_min, width, height)像素值
                        x_center = center_x_norm * width
                        y_center = center_y_norm * height
                        box_width = width_norm * width
                        box_height = height_norm * height

                        x_min = x_center - (box_width / 2)
                        y_min = y_center - (box_height / 2)

                        # COCO标注：x_min, y_min, bbox_width, bbox_height
                        bbox_coco = [float(x_min), float(y_min), float(box_width), float(box_height)]

                        # 确保边界框有效 (非负，非零宽度/高度)
                        if box_width <= 0 or box_height <= 0:
                            print(f"警告: 无效边界框尺寸在 {label_path}，跳过。")
                            continue

                        # 添加标注信息到COCO JSON
                        split_coco_data["annotations"].append({
                            "id": annotation_id_counter,
                            "image_id": current_image_id,
                            "category_id": class_id,
                            "bbox": bbox_coco,
                            "area": float(box_width * box_height),
                            "iscrowd": 0 # 对于标准目标检测通常设为0
                        })
                        annotation_id_counter += 1
            else:
                print(f"警告: 未找到标注文件 {label_path}，图片 {img_name} 将没有标注。")

        # 将当前split的COCO数据写入JSON文件
        output_json_path = os.path.join(coco_annotations_dir, f'instances_{split}.json')
        with open(output_json_path, 'w') as f:
            json.dump(split_coco_data, f, indent=4)
        print(f"{split} 阶段的COCO JSON已保存到: {output_json_path}")
        print(f"已处理 {len(split_coco_data['images'])} 张图片和 {len(split_coco_data['annotations'])} 条标注用于 {split} 阶段。")

    print(f"\n转换完成！COCO格式数据集已保存到: {coco_output_root}")

# --- 如何使用 ---
if __name__ == "__main__":
    # 请根据你的实际情况修改这些路径和类别名称
    YOLO_DATASET_ROOT = 'datasets_yolo' # 你的YOLO数据集的根目录
    COCO_OUTPUT_ROOT = 'datasets_coco'     # 转换后COCO数据集的输出目录

    # 你的40个垃圾类别名称，顺序必须与YOLO标注中的class_id (0-39) 对应
    # 示例：
    CLASS_NAMES = [
        'Paper', 
        'Cardboard', 
        'Glass', 
        'Plastic', 
        'Metal', 
        'Zip-topcan', 
        'Glassbottles', 
        'Batteries', 
        'Fluorescentlighttubes', 
        'Paintsandsolvents', 
        'Chemicals', 
        'Medications', 
        'Foodscraps', 
        'Vegetablepeels', 
        'Fruitpeels', 
        'Coffeegrounds', 
        'Tealeaves',
        'Planttrimmings', 
        'SoiledPlastic',
        'TornTextiles',
        'Ceramics',
        'Cigarettebutts',
        'Wood',
        'Drywall', 
        'Bricks', 
        'Concrete', 
        'Furniture', 
        'Appliance', 
        'Mattresses', 
        'Computer', 
        'Mobilephones', 
        'Televisions', 
        'Blood-soakedgauze', 
        'Needles', 
        'MedicalWaste', 
        'OldCothes', 
        'Bedding', 
        'Oldbook', 
        'Foodbox', 
        'Plasticbag'
    ]
    # 如果你有实际的类别文件（例如 obj.names），你可以这样加载：
    # with open(os.path.join(YOLO_DATASET_ROOT, 'data.yaml'), 'r') as f:
    #     data = yaml.safe_load(f)
    # CLASS_NAMES = data['names']

    # 调用转换函数
    yolo_to_coco(YOLO_DATASET_ROOT, COCO_OUTPUT_ROOT, CLASS_NAMES)

    # 验证输出目录结构
    print(f"\n输出目录结构示例:")
    print(f"{COCO_OUTPUT_ROOT}/")
    print(f"├─images/")
    print(f"│  ├─train/")
    print(f"│  ├─val/")
    print(f"│  └─test/")
    print(f"└─annotations/")
    print(f"   ├─instances_train.json")
    print(f"   ├─instances_val.json")
    print(f"   └─instances_test.json")